import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy

dt = 24 * 60 * 60
dt_times = 4500
precision_default = 1000

G = 6.67408e-11
AU = 1.495978707e11
M = {
    "Sun": 1.988e30,
    "Earth": 5.9722e24,
    "Jupyter": 1.8986e27
}
ORBIT = {
    "Earth":{
        "a": 1.496e11,
        "c": 1.4958e11
    },
    "Jupyter":{
        "a": (5.4588 + 4.9501) * AU / 2,
        "c": (5.4588 - 4.9501) * AU / 2
    }
}
PLANET = ["Earth", "Jupyter"]
POSITION = [
    ORBIT["Earth"]["a"] - ORBIT["Earth"]["c"],
    0,
    -(ORBIT["Jupyter"]["a"] + ORBIT["Jupyter"]["c"]),
    0
]
VELOCITY = [
    0,
    -np.sqrt(G * M["Sun"] * (2 * ORBIT["Earth"]["a"] - POSITION[0]) / (2 * ORBIT["Earth"]["a"] * POSITION[0])),
    0,
    np.sqrt(-G * M["Sun"] * (2 * ORBIT["Jupyter"]["a"] + POSITION[2]) / (2 * ORBIT["Jupyter"]["a"] * POSITION[2]))
]

def Simulation(position, velocity, t):
    x_e, y_e, x_j, y_j = deepcopy(position)
    vx_e, vy_e, vx_j, vy_j = deepcopy(velocity)
    r_ej = np.sqrt(np.square(x_e - x_j) + np.square(y_e - y_j))
    for i, x, y, vx, vy in zip([0, 1], [x_e, x_j], [y_e, y_j], [vx_e, vx_j], [vy_e, vy_j]):
        r = np.sqrt(np.square(x) + np.square(y))
        ax = -(x * np.power(r, -3) + np.power(-1, i) * (x_e - x_j) * np.power(r_ej, -3)) * G * M[PLANET[i]]
        ay = -(y * np.power(r, -3) + np.power(-1, i) * (y_e - y_j) * np.power(r_ej, -3)) * G * M[PLANET[i]]
        velocity[i * 2] = vx + ax * t
        velocity[i * 2 + 1] = vy + ay * t
        position[i * 2] = x + velocity[i * 2] * t
        position[i * 2 + 1] = y + velocity[i * 2 + 1] * t
    return velocity, position

if __name__ == "__main__":
    indexes = np.arange(0, dt_times * dt, dt)
    positions = np.array([POSITION])
    position = positions[-1]
    velocity = VELOCITY
    for i in indexes[:-1]:
        position, velocity = Simulation(position, velocity, dt)
        positions = np.append(positions, position)
    positions = positions.T
    plt.plot(positions[0], positions[1])
    plt.plot(positions[2], positions[4])
    plt.show()